﻿/*=============================================================================
#       COPYRIGHT NOTICE
#       Copyright (c) 2020
#       All rights reserved
#
#       @author       :	songzhicheng
#       @file         :	NtDemoApp.cpp
#       @date         :	2020/07/01
#       @description  :	NT框架驱动测试程序
=============================================================================*/

#include <iostream>
#include <Windows.h>
#include <tchar.h>
#include "helper_driver.h"
#include "../NtDemoDrv/common.h"

#define DRIVER_NAME		_T("NtDemoDrv")
#define DRIVER_PATH		_T("C:\\NtDemoDrv\\NtDemoDrv.sys")

#define DEVICE_NAME		_T("\\\\.\\NtDemo_device")

DWORD DoRead();			//设备操作:Read
DWORD DoWrite();		//设备操作:Write
DWORD DoDeviceIO();		//设备操作:DeviceIo

//注意！！本测试程序涉及的服务API、文件API都需要在管理员权限下运行
int main()
{
	DWORD dwErr = 0;
	do
	{
		printf("按任意键开始驱动测试\n");
		system("PAUSE");

		dwErr = DriverInstall(DRIVER_NAME, DRIVER_PATH, SERVICE_DEMAND_START);
		if (dwErr)
			break;
		printf("驱动安装：ok\n");

		system("PAUSE");
		dwErr = DriverStart(DRIVER_NAME);
		if (dwErr)
			break;
		printf("驱动启动：ok\n");

		system("PAUSE");
		dwErr = DoRead();
		if (dwErr)
			printf("DoRead：失败：%d\n", dwErr);
		else
			printf("DoRead：ok\n");

		system("PAUSE");
		dwErr = DoWrite();
		if (dwErr)
			printf("DoWrite：失败：%d\n", dwErr);
		else
			printf("DoWrite：ok\n");

		system("PAUSE");
		dwErr = DoDeviceIO();
		if (dwErr)
			printf("DoDeviceIO：失败：%d\n", dwErr);
		else
			printf("DoDeviceIO：ok\n");

		system("PAUSE");
		dwErr = DriverStop(DRIVER_NAME);
		if (dwErr)
			break;
		printf("驱动停止：ok\n");

		system("PAUSE");
		dwErr = DriverUninstall(DRIVER_NAME);
		if (dwErr)
			break;
		printf("驱动卸载：ok\n");

	} while (FALSE);

	if (dwErr)
	{
		printf("error happens:%d\n", dwErr);
	}

	return dwErr;
}

DWORD DoRead()
{
	DWORD dwErr = 0;
	HANDLE hDevice = INVALID_HANDLE_VALUE;

	WCHAR readBuf[100] = { 0 };
	DWORD dwBytesRead = 0;
	do
	{
		hDevice = CreateFile(DEVICE_NAME,
			GENERIC_READ | GENERIC_WRITE,
			0,		// share mode none
			NULL,	// no security
			OPEN_EXISTING,
			FILE_ATTRIBUTE_NORMAL,
			NULL);		// no template
		if (hDevice == INVALID_HANDLE_VALUE)
		{
			dwErr = GetLastError();
			break;
		}
		printf("驱动设备打开：ok\n");

		if (!ReadFile(hDevice, readBuf, sizeof(readBuf), &dwBytesRead, NULL))
		{
			dwErr = GetLastError();
			printf("ReadFile fail: %d, dwBytesRead:%d\n", dwErr, dwBytesRead);
		}
		else
		{
			printf("Read data: %ws\n", readBuf);
		}

	} while (FALSE);

	if (hDevice != INVALID_HANDLE_VALUE)
		CloseHandle(hDevice);

	return dwErr;
}

DWORD DoWrite()
{
	DWORD dwErr = 0;
	HANDLE hDevice = INVALID_HANDLE_VALUE;

	WCHAR writeBuf[] = L"data to write";
	DWORD dwBytesWritten = 0;
	do
	{
		hDevice = CreateFile(DEVICE_NAME,
			GENERIC_READ | GENERIC_WRITE,
			0,		// share mode none
			NULL,	// no security
			OPEN_EXISTING,
			FILE_ATTRIBUTE_NORMAL,
			NULL);		// no template
		if (hDevice == INVALID_HANDLE_VALUE)
		{
			dwErr = GetLastError();
			break;
		}
		printf("驱动设备打开：ok\n");

		if (!WriteFile(hDevice, writeBuf, sizeof(writeBuf), &dwBytesWritten, NULL))
		{
			dwErr = GetLastError();
			printf("WriteFile fail: %d, dwBytesWritten:%d\n", dwErr, dwBytesWritten);
		}
		else
		{
			printf("WriteFile ok: %ws\n", writeBuf);
		}
	} while (FALSE);

	if (hDevice != INVALID_HANDLE_VALUE)
		CloseHandle(hDevice);

	return dwErr;
}

DWORD DoDeviceIO()
{
	DWORD dwErr = 0;
	HANDLE hDevice = INVALID_HANDLE_VALUE;

	WCHAR inBuf[] = L"iotest";
	DWORD inBufSize = sizeof(inBuf);
	WCHAR* outBuf = NULL;
	do
	{
		hDevice = CreateFile(DEVICE_NAME,
			GENERIC_READ | GENERIC_WRITE,
			0,		// share mode none
			NULL,	// no security
			OPEN_EXISTING,
			FILE_ATTRIBUTE_NORMAL,
			NULL);		// no template
		if (hDevice == INVALID_HANDLE_VALUE)
		{
			dwErr = GetLastError();
			break;
		}
		printf("驱动设备打开：ok\n");

		//If lpOverlapped is NULL, lpBytesReturned cannot be NULL.
		printf("IOCTRL_TEST1\n");
		DWORD dwBytesNeeded = 0;
		if (!DeviceIoControl(hDevice, IOCTL_TEST1, inBuf, inBufSize,
			NULL, 0, &dwBytesNeeded, NULL))
		{
			dwErr = GetLastError();
			printf("IOCTRL_TEST1 fail: %d, dwBytesNeeded:%d\n", dwErr, dwBytesNeeded);
		}

		printf("IOCTRL_TEST2\n");
		dwBytesNeeded = 0;
		if (!DeviceIoControl(hDevice, IOCTL_TEST2, inBuf, inBufSize,
			NULL, 0, &dwBytesNeeded, NULL))
		{
			dwErr = GetLastError();
			printf("IOCTRL_TEST2 fail: %d, dwBytesNeeded:%d\n", dwErr, dwBytesNeeded);
			if (dwErr == ERROR_MORE_DATA)
			{
				outBuf = (WCHAR*)malloc(dwBytesNeeded);
				if (outBuf)
				{
					if (!DeviceIoControl(hDevice, IOCTL_TEST2, inBuf, inBufSize,
						outBuf, dwBytesNeeded, &dwBytesNeeded, NULL))
					{
						dwErr = GetLastError();
						printf("IOCTRL_TEST2 fail again: %d, dwBytesNeeded:%d\n", dwErr, dwBytesNeeded);
					}
					else
					{
						printf("IOCTRL_TEST2 ok: %ws\n", outBuf);
						dwErr = 0;
					}
					free(outBuf);
					outBuf = NULL;
				}
			}
		}

	} while (FALSE);

	if(hDevice != INVALID_HANDLE_VALUE)
		CloseHandle(hDevice);

	return dwErr;
}